#include "vio/front_end/patch_matcher.h"

namespace vio {
namespace warp {

void GetWarpMatrixAffine(const Camera::Ptr cam_ref,
                         const Camera::Ptr cam_cur,
                         const Eigen::Vector2d& px_ref,
                         const Eigen::Vector3d& f_ref,
                         const double& depth_ref,
                         const Sophus::SE3d& T_cur_ref,
                         const int level_ref,
                         Eigen::Matrix2d& A_cur_ref) {
  // compute affine warp matrix A_ref_cur
  const int half_patch_size = 5;
  // 3D point in camera frame
  const Eigen::Vector3d xyz_ref(f_ref * depth_ref);
  Eigen::Vector3d xyz_du_ref(cam_ref->pixel2camera(
      px_ref + Eigen::Vector2d(half_patch_size, 0) * (1 << level_ref)));
  Eigen::Vector3d xyz_dv_ref(cam_ref->pixel2camera(
      px_ref + Eigen::Vector2d(0, half_patch_size) * (1 << level_ref)));
  // locally planar patch with same depth
  xyz_du_ref *= xyz_ref[2] / xyz_du_ref[2];
  xyz_dv_ref *= xyz_ref[2] / xyz_dv_ref[2];
  // reproject 3 pixels from ref pixel frame to cur pixel frame
  const Eigen::Vector2d px_cur(cam_cur->camera2pixel(T_cur_ref * xyz_ref));
  const Eigen::Vector2d px_du(cam_cur->camera2pixel(T_cur_ref * xyz_du_ref));
  const Eigen::Vector2d px_dv(cam_cur->camera2pixel(T_cur_ref * xyz_dv_ref));
  A_cur_ref.col(0) = (px_du - px_cur) / half_patch_size;
  A_cur_ref.col(1) = (px_dv - px_cur) / half_patch_size;
}

int GetBestSearchLevel(const Eigen::Matrix2d& A_cur_ref, const int max_level) {
  // compute patch level in other image
  int search_level = 0;
  double D = A_cur_ref.determinant();
  while (D > 3.0 && search_level < max_level) {
    search_level += 1;
    D *= 0.25;
  }
  return search_level;
}

void WarpAffine(const Eigen::Matrix2d& A_cur_ref,
                const cv::Mat& img_ref,
                const Eigen::Vector2d& px_ref,
                const int level_ref,
                const int search_level,
                const int half_patch_size,
                uint8_t* patch_ptr) {
  const int patch_size = half_patch_size * 2;
  const Eigen::Matrix2f A_ref_cur = A_cur_ref.inverse().cast<float>();
  if (std::isnan(A_ref_cur(0, 0))) {
    LOG(INFO) << "Affine warp is NaN, probably camera has no translation";
    return;
  }

  // perform the warp on a larger patch
  const Eigen::Vector2f px_ref_pyr = px_ref.cast<float>() / (1 << level_ref);
  for (int y = 0; y < patch_size; y++) {
    for (int x = 0; x < patch_size; x++, patch_ptr++) {
      Eigen::Vector2f px_patch(x - half_patch_size, y - half_patch_size);
      px_patch *= (1 << search_level);
      const Eigen::Vector2f px(A_ref_cur * px_patch + px_ref_pyr);
      if (px[0] < 0 || px[1] < 0 || px[0] >= img_ref.cols - 1 ||
          px[1] >= img_ref.rows - 1) {
        *patch_ptr = 0;
      } else {
        *patch_ptr = vision_utils::InterpolateMat_8u(img_ref, px[0], px[1]);
      }
    }
  }
}
}  // namespace warp

// A = [R*x1, x2]
// A * [s1, s2]' = t
bool DepthFromTriangulation(const Sophus::SE3d& T_search_ref,
                            const Eigen::Vector3d& f_ref,
                            const Eigen::Vector3d& f_cur,
                            double& depth) {
  Eigen::Matrix<double, 3, 2> A;
  A << T_search_ref.rotationMatrix() * f_ref, f_cur;
  const Eigen::Matrix2d AtA = A.transpose() * A;
  if (AtA.determinant() < 0.000001) {
    return false;
  }
  const Eigen::Vector2d depth2 =
      -AtA.inverse() * A.transpose() * T_search_ref.translation();
  depth = std::abs(depth2[0]);
  return true;
}

// get center part of patch with border
void PatchMatcher::CreatePatchFromPatchWithBorder() {
  uint8_t* ref_patch_ptr = patch_;
  for (int y = 1; y < patch_size_ + 1; y++, ref_patch_ptr += patch_size_) {
    uint8_t* ref_patch_border_ptr =
        patch_with_border_ + y * (patch_size_ + 2) + 1;
    for (int x = 0; x < patch_size_; x++) {
      ref_patch_ptr[x] = ref_patch_border_ptr[x];
    }
  }
}

bool PatchMatcher::FindMatchDirect(const Point& point,
                                   const Frame& cur_frame,
                                   Eigen::Vector2d& px_cur) {
  if (!point.GetCloseViewObs(cur_frame.pos(), ref_ftr_)) {
    return false;
  }

  auto ftr_f_ptr = ref_ftr_->frame_ptr.lock();
  if (ftr_f_ptr && ftr_f_ptr->cam_ptr_->IsInFrame(
                       ref_ftr_->px.cast<int>() / (1 << ref_ftr_->level),
                       half_patch_size_ + 2,
                       ref_ftr_->level)) {
    return false;
  }

  // warp affine
  warp::GetWarpMatrixAffine(ftr_f_ptr->cam_ptr_,
                            cur_frame.cam_ptr_,
                            ref_ftr_->px,
                            ref_ftr_->bearing,
                            (ftr_f_ptr->pos() - point.pos_).norm(),
                            cur_frame.T_f_w_ * ftr_f_ptr->T_f_w_.inverse(),
                            ref_ftr_->level,
                            A_cur_ref_);
  search_level_ = warp::GetBestSearchLevel(A_cur_ref_, 2);
  warp::WarpAffine(A_cur_ref_,
                   ftr_f_ptr->img_pyr_[ref_ftr_->level],
                   ref_ftr_->px,
                   ref_ftr_->level,
                   search_level_,
                   half_patch_size_ + 1,
                   patch_with_border_);
  CreatePatchFromPatchWithBorder();

  // px_cur should be set
  Eigen::Vector2d px_scaled(px_cur / (1 << search_level_));

  bool success = false;
  if (ref_ftr_->type == Feature::EDGELET) {
    Eigen::Vector2d dir_cur(A_cur_ref_ * ref_ftr_->grad);
    dir_cur.normalize();
    success = feature_alignment::Align1D(cur_frame.img_pyr_[search_level_],
                                         dir_cur.cast<float>(),
                                         patch_with_border_,
                                         patch_,
                                         options_.align_max_iter,
                                         px_scaled,
                                         h_inv_);

  } else {
    success = feature_alignment::Align2D(cur_frame.img_pyr_[search_level_],
                                         patch_with_border_,
                                         patch_,
                                         options_.align_max_iter,
                                         px_scaled);
  }
  px_cur = px_scaled * (1 << search_level_);
  return success;
}

bool PatchMatcher::FindEpipolarMatchDirect(const Frame& ref_frame,
                                           const Frame& cur_frame,
                                           const Feature& ref_ftr,
                                           const double& d_estimate,
                                           const double& d_min,
                                           const double& d_max,
                                           double& depth) {
  Sophus::SE3d T_cur_ref = cur_frame.T_f_w_ * ref_frame.T_f_w_.inverse();
  int zmssd_best = PatchScore::Threshold();
  Eigen::Vector2d uv_best;

  // compute start adn end of epipolar line in old_kf for match search, on unit
  // plane!
  Eigen::Vector2d A =
      math_utils::project2d(T_cur_ref * (ref_ftr.bearing * d_min));
  Eigen::Vector2d B =
      math_utils::project2d(T_cur_ref * (ref_ftr.bearing * d_max));

  // compute affine warp matrix
  warp::GetWarpMatrixAffine(ref_frame.cam_ptr_,
                            cur_frame.cam_ptr_,
                            ref_ftr.px,
                            ref_ftr.bearing,
                            d_estimate,
                            T_cur_ref,
                            ref_ftr.level,
                            A_cur_ref_);

  // feature pre-selection
  reject_ = false;
  if (ref_ftr.type == Feature::EDGELET &&
      options_.epi_search_edgelet_filtering) {
    const Eigen::Vector2d grad_cur = (A_cur_ref_ * ref_ftr.grad).normalized();
    const double cos_angle = std::fabs(grad_cur.dot(epi_dir_.normalized()));
    // grad and epi should be have same direction
    if (cos_angle < options_.epi_search_edgelet_max_angle) {
      reject_ = true;
      return false;
    }
  }

  search_level_ = warp::GetBestSearchLevel(A_cur_ref_, 2);

  // find length of search range on epipolar line
  Eigen::Vector2d px_A(
      cur_frame.cam_ptr_->camera2pixel(math_utils::unproject2d(A)));
  Eigen::Vector2d px_B(
      cur_frame.cam_ptr_->camera2pixel(math_utils::unproject2d(B)));
  epi_length_ = (px_A - px_B).norm() / (1 << search_level_);

  // warp reference patch at ref_level
  warp::WarpAffine(A_cur_ref_,
                   ref_frame.img_pyr_[ref_ftr.level],
                   ref_ftr.px,
                   ref_ftr.level,
                   search_level_,
                   half_patch_size_ + 1,
                   patch_with_border_);
  CreatePatchFromPatchWithBorder();

  if (epi_length_ < 2.0) {
    px_cur_ = (px_A + px_B) / 2.0;
    Eigen::Vector2d px_scaled(px_cur_ / (1 << search_level_));
    bool res;
    if (options_.align_1d) {
      res = feature_alignment::Align1D(cur_frame.img_pyr_[search_level_],
                                       (px_A - px_B).cast<float>().normalized(),
                                       patch_with_border_,
                                       patch_,
                                       options_.align_max_iter,
                                       px_scaled,
                                       h_inv_);
    } else {
      res = feature_alignment::Align2D(cur_frame.img_pyr_[search_level_],
                                       patch_with_border_,
                                       patch_,
                                       options_.align_max_iter,
                                       px_scaled);
    }
    if (res) {
      px_cur_ = px_scaled * (1 << search_level_);
      if (DepthFromTriangulation(T_cur_ref,
                                 ref_ftr.bearing,
                                 cur_frame.cam_ptr_->pixel2camera(px_cur_),
                                 depth)) {
        return true;
      }
    }
    return false;
  }

  // one step per pixel
  size_t n_steps = epi_length_ / 0.7;
  Eigen::Vector2d step = epi_dir_ / n_steps;

  if (n_steps > options_.max_epi_search_steps) {
    LOG(INFO) << "skip epipolar search: " << n_steps
              << " evaluations, px_lenght: " << epi_length_
              << ", d_min: " << d_min << ", d_max: " << d_max;
    return false;
  }

  // for matching, precompute sum and sum2 of warped reference patch
  int pixel_sum = 0;
  int pixel_sum_square = 0;
  PatchScore patch_score(patch_);

  // now we sample along the epipolar line
  Eigen::Vector2d uv = B - step;
  Eigen::Vector2i last_checked_pxi(0, 0);
  ++n_steps;
  for (size_t i = 0; i < n_steps; i++, uv += step) {
    Eigen::Vector2d px(
        cur_frame.cam_ptr_->camera2pixel(math_utils::unproject2d(uv)));
    // +0.5 to round to closest int
    Eigen::Vector2i pxi(px[0] / (1 << search_level_) + 0.5,
                        px[1] / (1 << search_level_) + 0.5);
    if (pxi == last_checked_pxi) {
      continue;
    }
    // TODO interpolation would probably be a good idea
    uint8_t* cur_patch_ptr =
        cur_frame.img_pyr_[search_level_].data +
        (pxi[1] - half_patch_size_) * cur_frame.img_pyr_[search_level_].cols +
        (pxi[0] - half_patch_size_);
    int zmssd = patch_score.ComputeScore(
        cur_patch_ptr, cur_frame.img_pyr_[search_level_].cols);

    if (zmssd < zmssd_best) {
      zmssd_best = zmssd;
      uv_best = uv;
    }
  }

  if (zmssd_best < PatchScore::Threshold()) {
    if (options_.subpix_refinement) {
      px_cur_ = cur_frame.cam_ptr_->camera2pixel(math_utils::unproject2d(uv));
      Eigen::Vector2d px_scaled(px_cur_ / (1 << search_level_));
      bool res;
      if (options_.align_1d) {
        res =
            feature_alignment::Align1D(cur_frame.img_pyr_[search_level_],
                                       (px_A - px_B).cast<float>().normalized(),
                                       patch_with_border_,
                                       patch_,
                                       options_.align_max_iter,
                                       px_scaled,
                                       h_inv_);
      } else {
        res = feature_alignment::Align2D(cur_frame.img_pyr_[search_level_],
                                         patch_with_border_,
                                         patch_,
                                         options_.align_max_iter,
                                         px_scaled);
      }
      if (res) {
        px_cur_ = px_scaled * (1 << search_level_);
        if (DepthFromTriangulation(T_cur_ref,
                                   ref_ftr.bearing,
                                   cur_frame.cam_ptr_->pixel2camera(px_cur_),
                                   depth)) {
          return true;
        }
      }
      return false;
    }
    px_cur_ =
        cur_frame.cam_ptr_->camera2pixel(math_utils::unproject2d(uv_best));
    if (DepthFromTriangulation(T_cur_ref,
                               ref_ftr.bearing,
                               math_utils::unproject2d(uv_best).normalized(),
                               depth)) {
      return true;
    }
  }
  return false;
}
}  // namespace vio
